import base64
from collections.abc import Generator
import os
import random
import sys
import codecs
import time
import traceback
from aiohttp_retry import Union
from sqlmodel import SQLModel
from termcolor import colored

import json
from openai.types.chat.chat_completion_message import ChatCompletionMessage
from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall, Function
import warnings
from operation import AGENT_TOOLS,router,screenAndXML,adbIme,defaultIme
from LLMservice import LLMService,PrettyMessages
warnings.filterwarnings("ignore", category=DeprecationWarning)
#https://platform.openai.com/docs/overview

class UIAgent:
    def __init__(self):
        self.messages = []
        self.LLMservice = LLMService()
        self.numChackMain=0
        self.curTask=""
    def stepPre(self, task: str) -> None:
        # 如下是现有文件目录信息以及根据任务的关键词自动搜索到的相关结果\n{ragstr}.
        self.curTask=task
        self.messages.clear()
        self.numChackMain=0
        self.messages.append({
            "role":"system",
            "content":f'''你是一个手机UI操作助手。每次用户都会发来屏幕截图和a11yTree信息。
你可以一次性调用多个函数，它们会依次执行。
如果你不再调用任何函数，任务将会结束。
'''
# 如果你调用TaskComplete工具，任务将会结束。
# 在结束任务之前，你需要说'正在检查某现象以确认任务是否完成'，然后用户会再次给你屏幕截图和a11yTree信息。然后你要仔细查看确认任务是否完成，如果完成了，最后调用TaskComplete工具和speakParams工具（告诉用户执行完毕）。如果没完成，继续操作手机。
        })
        print(PrettyMessages(self.messages))#system
        self.messages.append({
            "role":"user",
            "content":f"Now call functions. Remember, the task is {self.curTask}"
        })
        print(PrettyMessages(self.messages))#user
        sys.stdout.flush()
    def checkPre(self):
        self.messages=[self.messages[0]]
        print(PrettyMessages(self.messages))#system
        self.numChackMain+=1
        if self.numChackMain>1:
            return
        self.messages.append({
            "role":"user",
            "content":f"请通过截图仔细检查{self.curTask}是否完成。寻找任务被完整地完成的证据。如果已完成，就不需要再次调用函数了，以示任务完成。如果未完成，请继续操作手机。"
        })
        print(PrettyMessages(self.messages))#user
        sys.stdout.flush()
        return self.stepMain()
    def stepMain(self):
        resultdict = screenAndXML()
        reply = self.LLMservice.funcCallLLM(
            "THE FOLLOWING allyTree describtion may not cohere with the screenshot. Please use the screenshot as the main source of information.But if a11yTree's content is cohere with the screenshot,the coordinates had better be from a11y nodes if possible."+resultdict['a11y'],
            resultdict['path'],
            AGENT_TOOLS,
            self.messages
        )
        if reply.tool_calls:
            for tool_call in reply.tool_calls:
                name = tool_call.function.name
                args = json.loads(tool_call.function.arguments)
                tool_result = router(name, args)
                self.addPntToolRtn(tool_result,tool_call)
            return self.stepMain()
        return self.checkPre()
    def addPntToolRtn(self,tool_result:dict, tool_call: ChatCompletionMessageToolCall)->None:
        """
        # 注意，只会push toolcall result ，不会push toolcall到self.messages
        执行工具调用并构造响应消息
        
        Args:
            tool_name: 工具名称 
            tool_args: 工具参数
            tool_id: 工具调用ID
        Returns:
            构造的工具响应消息
        """
        tool_result["tool_name"] = tool_call.function.name
        result_to_send = tool_result.copy()
        if "a11y" in result_to_send:
            del result_to_send["a11y"]
        tool_message = {
            "role": "tool",
            "tool_call_id": tool_call.id, 
            "content": json.dumps(result_to_send)
        }
        self.messages.append(tool_message)
        print(PrettyMessages(self.messages))#toolreturn
        sys.stdout.flush()
        # if tool_call.function.name=="ScreenshotParams" and tool_result.get('success'):
        #     with open(tool_result['path'], 'rb') as image_file:
        #         imgstr=base64.b64encode(image_file.read()).decode('utf-8')
        #     img_message = {
        #         "role": "user",
        #         "content": [
        # {'type':'text','text':"THE FOLLOWING allyTree describtion may not cohere with the screenshot. Please use the screenshot as the main source of information."},
        # {"type":"text","text":tool_result['a11y']},
        # {"type":"image_url","image_url":{"url":f"data:image/png;base64,{imgstr}"}}]
        #     }
        #     self.messages.append(img_message)
        #     print(PrettyMessages(self.messages))#toolreturn
        #     sys.stdout.flush()
        return 
    def main(self,task:str):
        adbIme()
        self.stepPre(task)
        self.stepMain()
        defaultIme()
if __name__=="__main__":
    agent = UIAgent()
    agent.main("converse with the contact and send message")
    #AND MAKE SURE YOU PRESSED SEND AND SEE THAT THE MESSAGE IS ACTUALLY SENT,WITH GREEN BUBBLE ON THE RIGHT, BEFORE STOPPING
    while True:
        task=input("Enter the task: ")
        agent.main(task)
